from PyQt4 import QtGui, QtCore

import logging

import sys
import math
from pymunk import Vec2d
from State import *

def RunVisualize(world, controller):
    app = QtGui.QApplication(sys.argv)
    widget = Visualize(world, controller)
    widget.show()
    app.exec_()


class Visualize(QtGui.QWidget):
    def __init__(self, world, controller=None, parent=None):
        QtGui.QWidget.__init__(self, parent)
        self.setGeometry(50, 20, 600, 300)
        self.setWindowTitle('flexJoint')
        
        grid = QtGui.QGridLayout()
        
        label = QtGui.QLabel(self)
        fg = "QLabel {color:white}"
        label.setStyleSheet( fg )
        label.setMinimumWidth(200)

        self.worldWidget = VisualizeWorld(world, controller, self)
        self.worldWidget.setFocus()
        grid.addWidget(self.worldWidget, 0, 0, 1, 6)

        self.phaseSpaceWidget = VisualizePhaseSpace(world, controller, self)
        grid.addWidget(self.phaseSpaceWidget, 0, 6, 1, 6)
        
        # Make sure phase space is updated
        self.worldWidget.otherWidgetsToUpdate = [self.phaseSpaceWidget]
        
#        vbox = QtGui.QVBoxLayout()
#        vbox.addWidget(label)
#        vbox.addStretch(1)
#        grid.addLayout(vbox, 0, 6, 1, 2)

#        self.worldWidget.label = label 
		
        self.setLayout(grid)

        # set background color
        color = QtGui.QColor(0, 0, 0)
        palette = QtGui.QPalette()
        palette.setColor(QtGui.QPalette.Window, color)
        palette.setColor(QtGui.QPalette.Background, color)
        self.setPalette(palette)
        self.setBackgroundRole(QtGui.QPalette.Window)
        
class VisualizePhaseSpace(QtGui.QWidget):
    def __init__(self, world, controller, parent=None):
        QtGui.QWidget.__init__(self, parent)

        self.xMin = -2.0
        self.xMax = 2.0
        self.yMin = -30
        self.yMax = 30
        self.world = world       
        self.controller = controller 

    def paintAxes(self, paint):
        paint.setPen(QtGui.QColor(0, 0, 255))
        self.paintLine(paint, 0, self.yMin, 0, self.yMax)
        self.paintLine(paint, self.xMin, 0, self.xMax, 0)
        #paint.drawLine(2, self.sizeY-6, self.sizeX-3, self.sizeY-6)
        #paint.drawLine(5, 2, 5, self.sizeY-3)

    def toPaint(self, x, y):
        sizeX = self.xMax - self.xMin
        sizeY = self.yMax - self.yMin
        x = x - self.xMin
        y = y - self.yMin
        pX = self.sizeX * x/sizeX
        pY = self.sizeY - self.sizeY * y/sizeY
        return (pX, pY)

    def paintLine(self, paint, x1, y1, x2, y2):
        (pX1, pY1) = self.toPaint(x1, y1)
        (pX2, pY2) = self.toPaint(x2, y2)
        paint.drawLine(pX1, pY1, pX2, pY2)
    
    def paintCircle(self, paint, x, y, r):
        (pX, pY) = self.toPaint(x, y)
        paint.drawEllipse(pX-r, pY-r, 2*r, 2*r)
    
    def paintHistory(self, paint):
        events = self.controller.history.events
        if len(events) < 1:
            return
        sp = events[0].state
        paint.setPen(QtGui.QColor(0, 150, 0))
        for e in events[1:]:
            (x1, y1) = (sp.arm1.a, sp.arm1.aDot)
            s = e.state
            (x2, y2) = (s.arm1.a, s.arm1.aDot)
            self.paintLine(paint, x1, y1, x2, y2)
            sp = s
    
    def paintGoal(self, paint):
        paint.setPen(QtGui.QColor(255, 255, 255))
        if self.controller.currentGoalState != None:
            state = self.controller.currentGoalState
            (x, y) = (state.arm1.a, state.arm1.aDot)
            self.paintCircle(paint, x, y, 5)
        
    def paintTargetTrajectory(self, paint):
        dt = 0.01
        state = self.controller.previousState
        if state == None or self.controller.pathGenerator.goalState == None:
            return
        (x1, y1) = (state.arm1.a, state.arm1.aDot)
        n = 0
        while n<100 and stateDistanceWithoutControls(state, self.controller.pathGenerator.goalState) > 0.2:
            state = self.controller.pathGenerator.getNextState(state, dt)
            (x2, y2) = (state.arm1.a, state.arm1.aDot)
            self.paintLine(paint, x1, y1, x2, y2)
            (x1, y1) = (x2, y2)
            n = n + 1
        
    def paintModelPositions(self, paint):
        paint.setPen(QtGui.QColor(0, 255, 0))
        for model in self.controller.models:
            for state in model.goodStates:
                (x, y) = (state.arm1.a, state.arm1.aDot)
                self.paintCircle(paint, x, y, 5)
    
    def paintEvent(self, event):
        self.sizeX = self.size().width()
        self.sizeY = self.size().height()
        paint = QtGui.QPainter()
        paint.begin(self)
        self.paintAxes(paint)
        self.paintHistory(paint)
        self.paintGoal(paint)
        self.paintTargetTrajectory(paint)
        #self.paintModelPositions(paint)
        paint.end()

class VisualizeWorld(QtGui.QWidget):
    def __init__(self, world, controller, parent=None):
        QtGui.QWidget.__init__(self, parent)

        self.world = world       
        self.controller = controller 
        self.label = None
        self.tSlider = None
        self.runMode = False
        self.timer = QtCore.QBasicTimer()
        
        self.dt = 0.01
        self.time = 0
    
    def lToPaintX(self, x):
        pX = self.sizeX * x/self.world.sizeX
        return pX

    def lToPaintY(self, y):
        pY = self.sizeY * y/self.world.sizeY
        return pY
    
    def toPaint(self, x, y):
        pX = self.sizeX * x/self.world.sizeX
        pY = self.sizeY - self.sizeY * y/self.world.sizeY
        return (pX, pY)
    
    def drawLineInWorld(self, paint, p1, p2):
        (x1, y1) = self.toPaint(p1.x, p1.y)
        (x2, y2) = self.toPaint(p2.x, p2.y)
        paint.drawLine(x1, y1, x2, y2)
    
    def paintSpring(self, paint, start, end, nTurns, height):
        angle = (end-start).angle
        up = ((end-start) / float(nTurns)) / 4 + Vec2d(0, height/2.0).rotated(angle)
        down = ((end-start) / float(nTurns)) / 4 + Vec2d(0, -height/2.0).rotated(angle)
        p1 = start
        for i in range(nTurns):
            p2 = p1 + up
            self.drawLineInWorld(paint, p1, p2)
            p1 = p2
            p2 = p1 + down + down 
            self.drawLineInWorld(paint, p1, p2)
            p1 = p2
            p2 = p1 + up
            self.drawLineInWorld(paint, p1, p2)
            p1 = p2
    
    def paintBorder(self, paint):
        paint.setPen(QtGui.QColor(0, 0, 255))
        paint.drawLine(0, 0, 0, self.sizeY-1)
        paint.drawLine(0, self.sizeY-1, self.sizeX-1, self.sizeY-1)
        paint.drawLine(self.sizeX-1, self.sizeY-1, self.sizeX-1, 0)
        paint.drawLine(self.sizeX-1, 0, 0, 0)
    
    def paintBalls(self, paint):
        paint.setPen(QtGui.QColor(255, 255, 255))
        for b in self.world.balls:
            (x1, y1) = self.toPaint(b.body.position.x, b.body.position.y)
            r1 = self.lToPaintX(b.radius)
            r2 = self.lToPaintY(b.radius)
            paint.drawEllipse(x1-r1, y1-r2, 2*r1, 2*r2)
            # Mark angle
            x2 = x1 + math.sin(b.body.angle)*r1
            y2 = y1 + math.cos(b.body.angle)*r2
            paint.drawLine(x1, y1, x2, y2)

        # Paint connections between balls
        paint.setPen(QtGui.QColor(0, 255, 0))
        for c in self.world.visualBallConnections:
            b1 = self.world.balls[c[0]]
            b2 = self.world.balls[c[1]]
            (x1, y1) = self.toPaint(b1.body.position.x, b1.body.position.y)
            (x2, y2) = self.toPaint(b2.body.position.x, b2.body.position.y)
            paint.drawLine(x1, y1, x2, y2)
    
    def paintSegments(self, paint):
        paint.setPen(QtGui.QColor(255, 255, 255))
        for s in self.world.segments:
            p1 = s.body.position + s.a.rotated(math.degrees(s.body.angle))
            p2 = s.body.position + s.b.rotated(math.degrees(s.body.angle))
            (x1,y1) = self.toPaint(p1.x, p1.y)
            (x2,y2) = self.toPaint(p2.x, p2.y)
            paint.drawLine(x1, y1, x2, y2)
    
    def paintArmControls(self, paint):
        if self.world.arm1Shape1 != None:
            paint.setPen(QtGui.QColor(255,0,0))
            p = self.world.arm1Shape1.body.position
            xDisplacement = 10
            # Upper control
            p1 = p + Vec2d(-xDisplacement,0)
            paint.setPen(QtGui.QColor(80,80,80))
            self.drawLineInWorld(paint, p1 + Vec2d(0,2), p1 + Vec2d(0,8))
            p1 = p + Vec2d(-self.world.arm1Control1-xDisplacement, 0)
            paint.setPen(QtGui.QColor(255,0,0))
            self.drawLineInWorld(paint, p1 + Vec2d(0,2), p1 + Vec2d(0,8))
            # Lower control
            p2 = p + Vec2d(-xDisplacement,0)
            paint.setPen(QtGui.QColor(80,80,80))
            self.drawLineInWorld(paint, p2 + Vec2d(0,-2), p2 + Vec2d(0,-8))
            p2 = p + Vec2d(-self.world.arm1Control2-xDisplacement, 0)
            paint.setPen(QtGui.QColor(255,0,0))
            self.drawLineInWorld(paint, p2 + Vec2d(0,-2), p2 + Vec2d(0,-8))
            # Springs
            paint.setPen(QtGui.QColor(200,200,200))
            angle = self.world.arm1Shape1.body.angle
            x = 5.0*math.sin(angle)
            y = 5.0*math.cos(angle)
            self.paintSpring(paint, p1 + Vec2d(0,5), p + Vec2d(-x,y), 3, 3)
            self.paintSpring(paint, p2 + Vec2d(0,-5), p + Vec2d(x,-y), 3, 3)

    def paintEvent(self, event):
        self.sizeX = self.size().width()
        self.sizeY = self.size().height()
        paint = QtGui.QPainter()
        paint.begin(self)
        self.paintBorder(paint)
        self.paintBalls(paint)
        self.paintSegments(paint)
        self.paintArmControls(paint)
        paint.end()
                
    def timerEvent(self, event):
        self.step()
        #self.updateLabel()
        self.update()
        self.updateOther()
        
    def updateOther(self):
        for w in self.otherWidgetsToUpdate:
            w.update()
        
    def keyPressEvent(self, ev):
        t = ev.text()
        if t == "S": 
            self.step()
        elif t == "r":
            if self.runMode:
                self.runMode = False
                self.timer.stop()
            else:
                self.runMode = True
                self.timer.start(80, self)
        elif t == "w":
            self.world.arm1Control1 = self.world.arm1Control1 - 0.2 
        elif t == "q":
            self.world.arm1Control1 = self.world.arm1Control1 + 0.2 
        elif t == "s":
            self.world.arm1Control2 = self.world.arm1Control2 - 0.2 
        elif t == "a":
            self.world.arm1Control2 = self.world.arm1Control2 + 0.2 
        elif t == "H":
            self.controller.history.save("history.pickle") 
        self.update()
        self.updateOther()
            
    def updateLabel(self):
        s = "Time = %i\n" % (self.controller.time)
        self.updateLabelMap(s, state)
    
    def mousePressEvent(self, ev):
        x = ev.pos().x()
        y = ev.pos().y()
        self.update()
        self.updateOther()

    def mouseReleaseEvent(self, ev):
        self.update()
        self.updateOther()

    def mouseMoveEvent(self, ev):
        #self.paintSelectedAgentDX = ev.pos().x() - self.mousePressPosition.x()
        #self.paintSelectedAgentDY = ev.pos().y() - self.mousePressPosition.y()
        self.update()
        self.updateOther()
        
    def step(self):
        if self.controller != None:
            state = self.world.getState()
            #print self.controller.active(state, endState)
            action = self.controller.act(state, self.time, self.dt)
            self.world.implement(action)
        self.time = self.time + self.dt
        self.world.step(self.dt)





